# 系统库
import time
import re
from typing import Dict
import json
### langchain 公用代码
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain.schema import SystemMessage
from langchain_community.chat_message_histories import SQLChatMessageHistory


# 自定义库引入
from common import config
from utils.upload_about import get_file_info,set_file_info,remove_local_file
from utils.upload_about import moonshot_get_file_content,moonshot_get_file_obj,upload_file_to_local


def get_session_history(session_id):
    return SQLChatMessageHistory(session_id, "sqlite:///memory/history.db")

# 创建llm模型
# section 是conf.ini里面的配置项
def init_model(section,cache=True,streaming=False) :
    # 模型定义   
    model_name = config[section]['MODEL_NAME']
    openai_api_key = config[section]['API_KEY']
    base_url = config[section]['BASE_PATH']
    return ChatOpenAI(streaming=streaming,model_name=model_name,temperature=0.2,openai_api_key=openai_api_key,base_url=base_url,cache=cache)



def set_chain(llm,param:Dict):    
    files = param.get('files')
    msg_type = param.get('type')
    msg = param.get('msg')
    files_prompt= []
    is_file_msg=False
    if msg_type!='goods_info':
        if files:
            tmp = set_files_prompt(files)
            if tmp.get('err'):
                return tmp
            else:
                # 获取到文件解析结果 重置消息
                files_prompt = tmp['files_prompt']
                filenames = tmp['filenames']
                msg = f'[{filenames}]{msg}'
                is_file_msg = True
    sys = set_goods_sys_prompt(msg) if msg_type=='goods_info' else "您是用户的智能管家，会尽可能在200字以内回答用户的问题,如果是文件不要告诉用户文件名和文件类型等信息"

    messages = [
        (
                "system",
                sys ,
            ),
            ("human", "{input}"),
        ] if msg_type=='goods_info' else [
            *files_prompt,
            (
                "system",
                sys ,
            ),
            MessagesPlaceholder(variable_name="history", n_messages=8),
            ("human", "{input}"),
        ]
       
    prompt = ChatPromptTemplate.from_messages(messages)
    chain = prompt | llm
    res = chain if msg_type=="goods_info" else RunnableWithMessageHistory(
            chain,
            get_session_history,
            input_messages_key="input",
            history_messages_key="history",
        )
    return {
        "chain":res,
        "msg":msg,
        "is_file_msg":is_file_msg
    }

def set_goods_sys_prompt(msg):   
    group = re.match(r'^(.+)的用途和使用方法$',msg)
    if group :
        goods_name = group.group(1)        
        return  f"""您是用户的智能管家，先判断{goods_name}是不是一个物品，如果不是回答“无”即可。如果是，按如下格式回答给用户：
        \n*** 用途 ***\n
            1. ....\n
            2. ....\n
            ....\n
            \n*** 使用方法***\n
            1. ....\n
            2. ....\n
            ....\n
        ，无需其他的文字做末尾总结。
        """
    

def set_files_prompt(files):
    files_prompt =[]
    err_msg = None
    # 一个文件出错则全部终止
    # 文件上传到本地获取文件信息    
    if files:
        # [{filename,finalname,open_path}]
        local_files_info = upload_file_to_local(files)
        _type = type(local_files_info)
        if _type==list:
            # 上传成功
            for item in local_files_info:
                name = item['finalname']
                exact_info= get_file_info(name)
                if exact_info:                    
                    files_prompt.append(SystemMessage(exact_info['exact_content']))
                else:                     
                    file_obj_res = moonshot_get_file_obj(item,open(item['open_path'], 'rb'))
                    if file_obj_res.status_code==200:                        
                        file_obj = json.loads(file_obj_res.text)
                        file_content_res = moonshot_get_file_content(file_obj['id'])
                        if file_content_res.status_code==200:
                            file_content = file_content_res.text
                            files_prompt.append(SystemMessage(file_content))
                            set_file_info(file_content,file_obj['id'],item['finalname'])
                        # 抽取文件内容失败
                        else:
                            print(f'抽取内容失败：{file_content_res.text}')
                            f_name = item['finalname']
                            err_msg=f'{f_name}文件解析失败:抽取内容失败'
                            remove_local_file(local_files_info)
                            break
                    # 文件上传抽取文件对象失败 格式等不正确
                    else:
                        # 删除本地文件
                        remove_local_file(local_files_info)
                        f_name = item['finalname']
                        err_msg=f'{f_name}文件解析失败：抽取文件对象失败'
                        break
            if err_msg:
                return {"err":err_msg,"status":201}
            else:
                return {
                    "files_prompt":files_prompt,
                    "filenames":','.join([d['finalname'] for d in local_files_info])
                }
        else:
            # 出错了
            return local_files_info
    else:
        return {"err":"没有文件需要解析","status":201}


def get_result(chain,msg):
    try:
        return chain.invoke({"input":msg})
    except Exception as e:
        print(e)
        err = '请求过于频繁，请1s后重试' if e.response.status_code == 429 else str(e)
        return {"err":err,"status":201}

                

    





